import fitz  # PyMuPDF
from typing import List, Tuple, Dict, Any
import re


class PDFTableMerger:
    def __init__(self, pdf_path: str):
        self.doc = fitz.open(pdf_path)
        self.pages = [self.doc.load_page(i) for i in range(len(self.doc))]
        print(f"已加载PDF文件，共 {len(self.pages)} 页")
        
        # 调试模式标志
        self.debug = True

    def extract_clip_regions(self, page_num: int) -> List[Tuple[float, float, float, float]]:
        """
        提取页面中的所有Clip区域和潜在表格边界
        返回边界框列表 [x0, y0, x1, y1]
        """
        clip_regions = []

        # 获取页面的绘图指令
        drawings = self.pages[page_num].get_drawings()

        if self.debug:
            print(f"第 {page_num+1} 页有 {len(drawings)} 个绘图对象")

        for drawing in drawings:
            if self.debug:
                print(f"绘图对象包含的键: {list(drawing.keys())}")
                
            # 检查是否有clip区域
            if "clip" in drawing and drawing["clip"]:
                # 获取clip路径的边界框
                clip_rect = drawing["rect"]
                clip_regions.append((clip_rect.x0, clip_rect.y0, clip_rect.x1, clip_rect.y1))
                if self.debug:
                    print(f"找到Clip区域: {clip_rect}")
            
            # 补充方法：如果没有clip区域，尝试从drawings中提取可能的表格边界
            elif "rect" in drawing:
                # 检查这个矩形是否可能是表格的一部分
                rect = drawing["rect"]
                # 排除过小的矩形
                if (rect.width > 10 and rect.height > 10):
                    if self.debug:
                        print(f"添加潜在表格边界: {rect}")
                    clip_regions.append((rect.x0, rect.y0, rect.x1, rect.y1))

        return clip_regions

    def find_table_clips_near_margin(self, page_num: int, margin_type: str = "bottom",
                                    threshold: float = 30.0) -> List[Tuple[float, float, float, float]]:
        """
        查找靠近页面边缘的Clip区域（可能是表格被截断的地方）
        margin_type: "bottom" 或 "top"
        threshold: 距离边缘的阈值（点）- 增加默认值以提高检测率
        """
        page_rect = self.pages[page_num].rect
        clip_regions = self.extract_clip_regions(page_num)
        candidate_clips = []

        if self.debug:
            print(f"{"*"*20} 第 {page_num+1} 页 {margin_type} 区域 {"*"*20}")
            print(f"页面尺寸: {page_rect}")
            print(f"找到 {len(clip_regions)} 个潜在区域")

        for clip in clip_regions:
            if self.debug:
                print(f"检查区域: {clip}")
                
            x0, y0, x1, y1 = clip

            if margin_type == "bottom":
                # 检查是否靠近页面底部
                distance_to_bottom = page_rect.y1 - y1
                if self.debug:
                    print(f"  距离底部: {distance_to_bottom:.2f}")
                if 0 <= distance_to_bottom <= threshold:
                    candidate_clips.append(clip)
                    if self.debug:
                        print(f"  ✓ 添加到底部候选区域")

            elif margin_type == "top":
                # 检查是否靠近页面顶部
                distance_to_top = y0 - page_rect.y0
                if self.debug:
                    print(f"  距离顶部: {distance_to_top:.2f}")
                if 0 <= distance_to_top <= threshold:
                    candidate_clips.append(clip)
                    if self.debug:
                        print(f"  ✓ 添加到顶部候选区域")

        if self.debug:
            print(f"找到 {len(candidate_clips)} 个边缘候选区域")

        return candidate_clips

    def are_clips_aligned(self, clip1: Tuple[float, float, float, float],
                         clip2: Tuple[float, float, float, float],
                         tolerance: float = 10.0) -> bool:
        """
        检查两个Clip区域是否在水平方向上对齐
        tolerance: 允许的水平位置偏差（点）- 增加默认值以提高检测率
        """
        x0_1, _, x1_1, _ = clip1
        x0_2, _, x1_2, _ = clip2

        # 计算偏差
        left_diff = abs(x0_1 - x0_2)
        right_diff = abs(x1_1 - x1_2)
        
        if self.debug:
            print(f"对齐检查 - 左边界偏差: {left_diff:.2f}, 右边界偏差: {right_diff:.2f}")
        
        # 检查左边界和右边界是否大致对齐
        left_aligned = left_diff <= tolerance
        right_aligned = right_diff <= tolerance
        
        # 额外检查：宽度是否相似
        width1 = x1_1 - x0_1
        width2 = x1_2 - x0_2
        width_ratio = min(width1, width2) / max(width1, width2, 1)
        width_similar = width_ratio > 0.8  # 宽度相差不超过20%
        
        if self.debug:
            print(f"宽度1: {width1:.2f}, 宽度2: {width2:.2f}, 宽度相似度: {width_ratio:.2f}")
        
        # 可以根据任一条件返回True，提高检测率
        return (left_aligned and right_aligned) or width_similar

    def find_cross_page_tables(self) -> List[Tuple[int, int]]:
        """
        查找跨页表格对
        返回包含跨页表格的页面索引对列表 [(page_i, page_j), ...]
        改进：添加文本连续性检查作为辅助方法
        """
        cross_page_tables = []

        for i in range(len(self.pages) - 1):
            if self.debug:
                print(f"\n检查第 {i+1} 页和第 {i+2} 页是否存在跨页表格")
            
            # 方法1: 使用原始的Clip区域对齐检测
            bottom_clips = self.find_table_clips_near_margin(i, "bottom")
            top_clips = self.find_table_clips_near_margin(i + 1, "top")

            found = False
            # 检查是否有对齐的Clip对
            for bottom_clip in bottom_clips:
                for top_clip in top_clips:
                    if self.are_clips_aligned(bottom_clip, top_clip):
                        cross_page_tables.append((i, i + 1))
                        if self.debug:
                            print(f"✓ 通过Clip区域对齐检测到跨页表格")
                        found = True
                        break  # 找到一对就跳出内层循环
                if found:
                    break
            
            # 如果方法1没找到，尝试方法2: 文本连续性检查
            if not found:
                if self._check_text_continuity(i, i+1):
                    cross_page_tables.append((i, i + 1))
                    if self.debug:
                        print(f"✓ 通过文本连续性检测到跨页表格")

        return cross_page_tables
        
    def _check_text_continuity(self, page1: int, page2: int) -> bool:
        """
        检查两页之间的文本连续性，辅助检测跨页表格
        """
        # 获取页面1底部和页面2顶部的文本
        page1_rect = self.pages[page1].rect
        page2_rect = self.pages[page2].rect
        
        # 定义底部和顶部区域
        bottom_region = (page1_rect.x0, page1_rect.y1 - 50, page1_rect.x1, page1_rect.y1)
        top_region = (page2_rect.x0, page2_rect.y0, page2_rect.x1, page2_rect.y0 + 50)
        
        # 提取文本
        bottom_text = self._extract_text_in_region(page1, bottom_region)
        top_text = self._extract_text_in_region(page2, top_region)
        
        if self.debug:
            print(f"文本连续性检查 - 第 {page1+1} 页底部文本: {bottom_text[:100]}...")
            print(f"文本连续性检查 - 第 {page2+1} 页顶部文本: {top_text[:100]}...")
        
        # 检查数字序列
        bottom_numbers = re.findall(r'\d+', bottom_text)
        top_numbers = re.findall(r'\d+', top_text)
        
        if bottom_numbers and top_numbers:
            try:
                last_bottom = int(bottom_numbers[-1])
                first_top = int(top_numbers[0])
                if first_top == last_bottom + 1:
                    if self.debug:
                        print(f"发现连续数字序列: {last_bottom} -> {first_top}")
                    return True
            except:
                pass
        
        # 检查文本相似性
        bottom_words = set(re.findall(r'\w+', bottom_text.lower()))
        top_words = set(re.findall(r'\w+', top_text.lower()))
        common_words = bottom_words & top_words
        
        if common_words and len(common_words) > 3:
            if self.debug:
                print(f"发现多个共同词汇: {common_words}")
            return True
        
        return False
        
    def _extract_text_in_region(self, page_num: int, region: Tuple[float, float, float, float]) -> str:
        """
        提取指定区域内的文本
        """
        x0, y0, x1, y1 = region
        
        # 获取页面文本，按块组织
        text_dict = self.pages[page_num].get_text("dict")
        
        texts = []
        for block in text_dict.get("blocks", []):
            if block["type"] == 0:  # 文本块
                bbox = block["bbox"]
                block_x0, block_y0, block_x1, block_y1 = bbox
                
                # 检查文本块是否在区域内或与区域重叠
                if (block_x0 < x1 and block_x1 > x0 and block_y0 < y1 and block_y1 > y0):
                    for line in block.get("lines", []):
                        for span in line.get("spans", []):
                            texts.append(span["text"])
        
        return " ".join(texts)

    def extract_table_text(self, page_num: int, clip_region: Tuple[float, float, float, float]) -> List[List[str]]:
        """
        提取指定Clip区域内的表格文本
        返回二维列表表示的行列数据
        改进：更好的行检测和文本提取
        """
        x0, y0, x1, y1 = clip_region

        if self.debug:
            print(f"提取第 {page_num+1} 页区域 {clip_region} 内的表格文本")

        # 获取页面文本，按块和行组织
        text_dict = self.pages[page_num].get_text("dict")

        table_data = []
        current_row = []
        last_y = None
        
        # 存储所有文本行
        text_lines = []

        for block in text_dict.get("blocks", []):
            if block["type"] == 0:  # 文本块
                bbox = block["bbox"]
                block_x0, block_y0, block_x1, block_y1 = bbox

                # 检查文本块是否在Clip区域内或与区域重叠
                if (block_x0 < x1 and block_x1 > x0 and block_y0 < y1 and block_y1 > y0):

                    for line in block.get("lines", []):
                        for span in line.get("spans", []):
                            text = span["text"].strip()
                            if text:
                                line_bbox = span["bbox"]
                                line_y = (line_bbox[1] + line_bbox[3]) / 2  # 文本行的中心Y坐标
                                text_lines.append((line_y, text, line_bbox))

        # 按Y坐标排序文本行
        text_lines.sort(key=lambda x: x[0])
        
        # 构建表格行
        for y, text, bbox in text_lines:
            # 检测是否是新行
            if last_y is None or abs(y - last_y) > 6:  # 稍微增加阈值，适应不同PDF的行间距
                if current_row:
                    table_data.append(current_row)
                    current_row = []
                last_y = y

            current_row.append(text)

        # 添加最后一行
        if current_row:
            table_data.append(current_row)

        if self.debug:
            print(f"提取到 {len(table_data)} 行表格数据")

        return table_data

    def merge_tables(self, page1: int, page2: int) -> List[List[str]]:
        """
        合并两个页面上的表格数据
        改进：处理可能的行重叠和更好的表格边界确定
        """
        # 获取两个页面上的Clip区域
        bottom_clips = self.find_table_clips_near_margin(page1, "bottom")
        top_clips = self.find_table_clips_near_margin(page2, "top")

        # 找到对齐的Clip对
        aligned_clips = []
        for bottom_clip in bottom_clips:
            for top_clip in top_clips:
                if self.are_clips_aligned(bottom_clip, top_clip):
                    aligned_clips.append((bottom_clip, top_clip))
                    if self.debug:
                        print(f"找到对齐的Clip对: 第 {page1+1} 页 {bottom_clip}, 第 {page2+1} 页 {top_clip}")
                    break

        if not aligned_clips:
            # 如果没有找到对齐的Clip对，尝试使用整个页面底部和顶部区域
            page1_rect = self.pages[page1].rect
            page2_rect = self.pages[page2].rect
            
            # 定义更大的区域
            bottom_region = (page1_rect.x0, page1_rect.y1 - 100, page1_rect.x1, page1_rect.y1)
            top_region = (page2_rect.x0, page2_rect.y0, page2_rect.x1, page2_rect.y0 + 100)
            
            if self.debug:
                print(f"未找到对齐的Clip对，尝试使用默认区域: {bottom_region}, {top_region}")
            
            aligned_clips = [(bottom_region, top_region)]

        # 提取两个表格部分的数据
        table_part1 = self.extract_table_text(page1, aligned_clips[0][0])
        table_part2 = self.extract_table_text(page2, aligned_clips[0][1])

        if self.debug:
            print(f"第 {page1+1} 页提取到 {len(table_part1)} 行数据")
            print(f"第 {page2+1} 页提取到 {len(table_part2)} 行数据")
        
        # 去重处理：检查是否有重叠的行
        if table_part1 and table_part2:
            # 检查最后几行和第一几行是否相似
            last_rows = table_part1[-2:] if len(table_part1) >= 2 else table_part1
            first_rows = table_part2[:2] if len(table_part2) >= 2 else table_part2
            
            # 简单的去重逻辑
            to_remove = 0
            for i, row1 in enumerate(reversed(last_rows)):
                for j, row2 in enumerate(first_rows):
                    # 如果两行内容相似，则认为是重复
                    if self._rows_are_similar(row1, row2):
                        to_remove = len(last_rows) - i
                        break
                if to_remove > 0:
                    break
            
            if to_remove > 0:
                if self.debug:
                    print(f"检测到 {to_remove} 行重叠内容，已去除重复")
                table_part1 = table_part1[:-to_remove]

        # 合并表格数据
        merged_table = table_part1 + table_part2

        return merged_table
        
    def _rows_are_similar(self, row1: List[str], row2: List[str]) -> bool:
        """
        检查两行是否相似（可能是重复行）
        """
        # 如果长度差异太大，肯定不相似
        if abs(len(row1) - len(row2)) > 1:
            return False
        
        # 计算相似单元格的比例
        min_len = min(len(row1), len(row2))
        similar_cells = 0
        
        for i in range(min_len):
            # 简单的文本相似度检查
            text1 = row1[i].strip().lower()
            text2 = row2[i].strip().lower()
            
            # 如果文本完全相同，或者一个是另一个的子集
            if text1 == text2 or text1 in text2 or text2 in text1:
                similar_cells += 1
        
        # 如果80%以上的单元格相似，则认为行相似
        return similar_cells / min_len >= 0.8

    def process_document(self):
        """
        处理整个文档，查找并合并所有跨页表格
        """
        # 查找所有跨页表格
        cross_page_pairs = self.find_cross_page_tables()

        # 存储合并后的表格
        merged_tables = {}

        for page1, page2 in cross_page_pairs:
            print(f"发现跨页表格: 第 {page1+1} 页和第 {page2+1} 页")

            # 合并表格
            merged_table = self.merge_tables(page1, page2)

            if merged_table:
                merged_tables[(page1, page2)] = merged_table
                print(f"成功合并表格，共 {len(merged_table)} 行")

                # 打印前几行作为示例
                for i, row in enumerate(merged_table[:3]):
                    print(f"行 {i+1}: {row}")
                if len(merged_table) > 3:
                    print("...")
            print()

        return merged_tables

# 使用示例
if __name__ == "__main__":
    import sys
    import os
    
    if len(sys.argv) < 2:
        print("用法: python flip-01.py <pdf文件路径>")
        print("示例: python flip-01.py 'D:\\工作\\02-文档智能\\05-测试数据\\01-跨页表格\\跨页表格01.pdf'")
        sys.exit(1)
    
    pdf_path = sys.argv[1]
    
    if not os.path.exists(pdf_path):
        print(f"错误: 找不到PDF文件 '{pdf_path}'")
        sys.exit(1)

    try:
        print(f"开始处理PDF文件: {pdf_path}")
        merger = PDFTableMerger(pdf_path)
        merged_tables = merger.process_document()

        if not merged_tables:
            print("未检测到跨页表格")
            print("\n可能的原因和解决方法:")
            print("1. PDF中的表格可能没有使用标准的Clip区域")
            print("2. 尝试使用我们提供的 improved_table_detection.py 脚本")
            print("3. 运行: python improved_table_detection.py '{pdf_path}'")

    except Exception as e:
        print(f"处理PDF时发生错误: {e}")
        import traceback
        traceback.print_exc()
